import numpy as np
import random
from tqdm import tqdm
import matplotlib.pyplot as plt


class GaussMultiClustering:
    def __init__(self, clusterCount, maxEpochs, path):
        self.path = path
        self.maxEpochs = maxEpochs
        self.readData()
        self.dimension = len(self.Data[0])
        self.n = len(self.Data)
        # 初始化模型参数
        self.k = clusterCount
        # 随机生成k个高斯混合系数
        self.alpha = np.random.uniform(size=self.k)
        self.alpha = self.alpha/sum(self.alpha)
        # 随机生成k个初始中心点（从所有点中随机选出k个）
        self.u = random.sample(list(self.Data), self.k)
        # 生成k个协方差矩阵,都是以0.5对角线
        self.β = np.empty(shape=(self.n, self.k))
        covAll = []
        for i in range(0, self.k):
            cov = np.diag([0.5]*self.dimension)
            covAll.append(cov)
        self.cov = np.array(covAll)

    def readData(self):
        res = []
        myfile = open(self.path, "r", encoding="utf-8")
        lines = myfile.readlines()
        for line in lines:
            x = [float(i) for i in line[:-1].split(',')]
            res.append(x)
        self.Data = np.array(res)

    def g(self, x, u, cov):
        cov_inv = np.linalg.pinv(cov)
        cov_det = np.linalg.det(cov)
        return np.exp(-1/2*((x-u).T.dot(cov_inv.dot(x-u))))/np.sqrt((2*np.pi)**self.dimension*abs(cov_det))

    def getβ(self):
        for i in range(0, self.n):
            for j in range(0, self.k):
                self.β[i, j] = self.alpha[j] * \
                    self.g(self.Data[i], self.u[j], self.cov[j])
            self.β[i] /= np.sum(self.β[i])

    def Run(self, SavePic):
        for ep in tqdm(range(1, self.maxEpochs+1)):
            self.getβ()
            for j in range(self.k):
                r_toal = np.sum(self.β[:, j])
                self.u[j] = np.sum([self.Data[i]*self.β[i, j]
                                    for i in range(self.n)], axis=0)/r_toal
                self.cov[j] = np.sum([self.β[i, j]*((self.Data[i]-self.u[j]).reshape((self.dimension, 1)).dot(
                    (self.Data[i]-self.u[j]).reshape((1, self.dimension)))) for i in range(self.n)], axis=0)/r_toal/self.dimension
                self.alpha[j] = r_toal/self.n
            self.makeC()
            if SavePic:
                self.draw(ep=ep, showPic=False)
        self.draw(0, True)

    def makeC(self):
        self.C = []
        for j in range(self.k):
            self.C.append([])
        for i in range(self.n):
            c_i = np.argmax(self.β[i, :])
            self.C[c_i].append(self.Data[i])

    def draw(self, ep, showPic):
        colors = ['green', 'blue', 'red', 'black', 'yellow',
                  'pink', 'brown', 'lightblue', 'orange']
        for i in range(len(self.C)):
            plt.scatter([d[0] for d in self.C[i]], [d[1]
                                                    for d in self.C[i]], color=colors[i], label=str(i))
        plt.scatter([d[0] for d in self.u], [d[1] for d in self.u],
                    color=colors[-1], marker='^', label='center')
        plt.xlabel('x')
        plt.ylabel('y')
        plt.legend()
        if showPic:
            plt.show()
        else:
            plt.savefig("./images/%d.jpg" % ep)
        plt.close("all")


if __name__ == "__main__":
    myGaussMulti = GaussMultiClustering(
        clusterCount=6, maxEpochs=200, path="./data1.txt")
    myGaussMulti.Run(SavePic=True)
